# -*- coding: utf-8 -*-
# 在ArcMap 10.8上运行通过，Python 2.7
# 并行化的要素转栅格（注意：本代码中不要包含矢量点连接csv文件的步骤，这样会很慢）
import os
import sys
import arcpy
from multiprocessing import Pool, cpu_count
import argparse

# 获取当前脚本所在目录的父目录
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)
from utils.utools import timer

def feature_to_raster(subset_params):
    fc_layer_name, value_field, output_raster_fullpath, cell_size, lyr_number = subset_params
    # 创建要素图层
    arcpy.MakeFeatureLayer_management(fc_layer_name, "temp_layer{0}".format(lyr_number))
    # 执行要素转栅格
    arcpy.AddMessage(output_raster_fullpath)
    if arcpy.Exists(output_raster_fullpath):
        arcpy.Delete_management(output_raster_fullpath)
    try:
        arcpy.conversion.FeatureToRaster("temp_layer{0}".format(lyr_number),
                                       value_field,
                                       output_raster_fullpath,
                                       cell_size)
    except arcpy.ExecuteError as e:
        arcpy.AddMessage(str(e))

@timer
def featureToRaster_parallel(gdb_path, fc_layer_name, value_field, output_tif_fullpath, cell_size, num_cpus=None):
    '''
    滑坡易发性中，将点要素转为指定分辨率的栅格图像

    :param gdb_path: GDB文件地理数据库
    :param fc_layer_name: GDB文件地理库中，点要素类的名称（请不要包含后缀）
    :param value_field: 需要转换为栅格的字段，一般是【滑坡发生的概率字段】
    :param output_tif_fullpath: 输出tif文件的完整路径，输出目录不能是GDB文件地理数据库，否则会有锁的问题
    :param cell_size: 输出tif文件的地理分辨率，一般是30（指的是30米分辨率）
    :param num_cpus: 工作的进程数目，默认值为None。如果为None，则num_cpus等于CPU总线程数-2
    :return: 无返回值
    '''
    arcpy.env.overwriteOutput = True
    arcpy.env.workspace = gdb_path
    # 使用os.path.dirname()获取文件所在的文件夹路径
    raster_out_dir = os.path.dirname(output_tif_fullpath)

    # 【Step 0】：确保输入的参数符合要求
    if not arcpy.Exists(fc_layer_name):
        arcpy.AddMessage("要素类不存在: {0}".format(fc_layer_name))
        exit()
    if fc_layer_name.endswith('.gpkg') or fc_layer_name.endswith('.shp'):
        arcpy.AddMessage("gpkg和shp文件格式不受支持！")
        exit()
    if not gdb_path.endswith('.gdb'):
        arcpy.AddMessage("工作区间必须是文件地理数据库，即gdb数据库！")
        exit()
    if raster_out_dir=='':
        arcpy.AddMessage("输出文件路径output_tif_fullpath必须是完整路径，不能是单独文件名！")
        exit()
    elif raster_out_dir.endswith('.gdb'):
        arcpy.AddMessage("输出的tif文件路径不能是在gdb数据库中，否则会产生锁的问题！！！")
        exit()
    if not output_tif_fullpath.endswith('.tif'):
        arcpy.AddMessage("输出的完整文件路径必须是tif格式，请确保文件名以.tif结束！！！")
        exit()
    # 使用os.path.exists()判断文件夹是否存在
    if not os.path.exists(raster_out_dir):
        os.makedirs(raster_out_dir)

    # 【Step 1】： 代码模式下，由于input_fc是要素类，需要转化为layer图层。在ArcGIS Pro中直接执行不需要
    # 由于management.AddJoin的第一个参数为图层，而不是要素类。因此先要通过MakeFeatureLayer_management方法创建一个要素图层。
    arcpy.MakeFeatureLayer_management(fc_layer_name,  fc_layer_name)

    # 【Step 2】： 确定将大的要素类分割为多少份，每份包含多少要素
    count = int(arcpy.GetCount_management(fc_layer_name).getOutput(0))
    if num_cpus is None or num_cpus > cpu_count() or num_cpus<=0:
        num_cpus = cpu_count() - 2  # 使用可用的CPU核心数
    num_parts = num_cpus
    features_each_parts = int(count / num_parts)
    arcpy.AddMessage("点的总数目count:{0}".format(count))
    arcpy.AddMessage("点文件被分割成num_parts:{0}份".format(num_parts))
    arcpy.AddMessage("每一份的点数目为features_each_parts:{0}".format(features_each_parts))


    # 【Step 3】： 将输入的大要素类分割为若干个小的要素类
    subset_params_list = []
    raster_list = []
    subfc_list =[]
    for i in range(num_parts):
        where_clause = "OBJECTID >= {} AND OBJECTID < {}".format(i * features_each_parts, (i + 1) * features_each_parts)
        subset_fc = os.path.join(gdb_path, "subset_{}".format(i))
        arcpy.SelectLayerByAttribute_management(fc_layer_name, "NEW_SELECTION", where_clause)
        arcpy.CopyFeatures_management(fc_layer_name, subset_fc)

        subset_raster_name = "{0}_{1}".format(os.path.basename(output_tif_fullpath)[:-4], i)
        subset_raster_fullpath = "{0}\\{1}.tif".format(raster_out_dir, subset_raster_name)

        subset_params_list.append((subset_fc, value_field, subset_raster_fullpath, cell_size, i))

        raster_list.append(subset_raster_fullpath)
        subfc_list.append(subset_fc)
    # print(subset_params_list)

    # 【Step 4】： 根据CPU数量，并行化执行要素转栅格
    pool = Pool(processes=num_cpus)  # 根据你的CPU核心数来设置进程数
    pool.map(feature_to_raster, subset_params_list)
    pool.close()
    pool.join()

    # 【Step 5】： 并行化结果为多个小栅格tif文件，将多个小栅格文件拼接为一个大的栅格
    arcpy.AddMessage('开始拼接......')
    output_merged_raster_name = os.path.basename(output_tif_fullpath)
    # 设置输出栅格的参数
    mosaic_type = "LAST"  # 拼接规则，LAST表示最后覆盖
    pixel_type = "32_BIT_FLOAT"  # 像素类型，根据实际情况选择
    number_of_bands = 1  # 栅格波段数，如果不确定，可以设置为1
    # 执行拼接
    arcpy.MosaicToNewRaster_management(
        input_rasters=raster_list,
        output_location=raster_out_dir,  # 输出到普通文件夹
        raster_dataset_name_with_extension=output_merged_raster_name,
        coordinate_system_for_the_raster="",  # 如果留空，则采用第一个输入栅格的空间参考
        pixel_type=pixel_type,
        cellsize=cell_size,
        number_of_bands=number_of_bands,
        mosaic_method=mosaic_type,
    )
    arcpy.AddMessage("拼接结果存储在：{0}".format(output_merged_raster_name))
    arcpy.AddMessage("栅格拼接完成！")

    # 【Step 6】： 删除中间文件
    for raster_path in raster_list:
        if arcpy.Exists(raster_path):
            arcpy.Delete_management(raster_path)
            arcpy.AddMessage("已删除栅格文件: {0}".format(raster_path))
        else:
            arcpy.AddMessage("文件不存在: {0}".format(raster_path))
    for feature_class_path in subfc_list:
        if arcpy.Exists(feature_class_path):
            arcpy.Delete_management(feature_class_path)
            arcpy.AddMessage("已删除要素类: {0}".format(feature_class_path))
        else:
            arcpy.AddMessage("文件不存在: {0}".format(feature_class_path))

def main():
    parser = argparse.ArgumentParser(description='要素转栅格并行化脚本')
    parser.add_argument('--gdb_path', required=True, help='GDB文件地理数据库的路径')
    parser.add_argument('--fc_layer_name', required=True, help='GDB文件地理库中，点要素类的名称（请不要包含后缀）')
    parser.add_argument('--value_field', required=True, help='需要转换为栅格的字段，一般是【滑坡发生的概率字段】')
    parser.add_argument('--output_tif_fullpath', required=True, help='输出tif文件的完整路径，输出目录不能是GDB文件地理数据库，否则会有锁的问题')
    parser.add_argument('--cell_size', type=int, required=True, help='输出tif文件的地理分辨率，一般是30（指的是30米分辨率）')
    parser.add_argument('--num_cpus', type=int, required=True, help='用于后台工作的CPU数量')

    args = parser.parse_args()
    
    # 调用要素转栅格的并行化版本
    featureToRaster_parallel(args.gdb_path, args.fc_layer_name, args.value_field, args.output_tif_fullpath, args.cell_size, args.num_cpus)

if __name__=="__main__":
    # gdb_path = r"C:\Users\Luoge7000P\Documents\ArcGIS\Projects\testFeatureToRaster\testFeatureToRaster.gdb"
    # fc_layer_name = "RasterUnits20220905joined"
    # value_field = "y_pred_prob"
    # output_tif_fullpath = "G:/tmpOut20211201/predRe20241029_115513_shp.tif"
    # cell_size = 30

    # 调用要素转栅格的并行化版本
    # featureToRaster_parallel(gdb_path, fc_layer_name, value_field, output_tif_fullpath, cell_size)
    main()
